#include "proxy_kpabe.h"
#include "kpabe_KPAbe.h"
#include "string.h"
#include <stdlib.h>
#include <stdio.h>
#include <iostream>
using namespace std;

int beta, alpha;
kpabe_pub_t* pub;
kpabe_msk_t* msk;
element_t* maskArray;
element_t* userDecKeys;

unordered_map<element_wrap,int, elementHash, equalElement> mapElement;
unordered_map<int,element_wrap> mapPlaintext; 
unordered_map<int, unsigned char*> userPolicyMap;

gint str_cmp(gconstpointer a, gconstpointer b){
	return strcmp((const char*)a,(const char*)b);
}

char** parseAttributeUniverse(char** argv, int size){
        GSList* alist;
        GSList* ap;
        alist = 0;
        ap = 0;
        for (int i=0; i<size; i++)
                parse_attribute_universe(&alist,argv[i]);
        alist = g_slist_sort(alist,str_cmp);
        int n = g_slist_length(alist);
        char** atts = (char**)malloc((n+1)*sizeof(char*));
        int i=0;
        for (ap = alist; ap; ap = ap->next)
                atts[i++] = (char*)ap->data;
        atts[i] = 0;
	return atts;
}

char** parseEncryptionAttributes(char** argv, int size){
	GSList* alist;
	GSList* ap;
	alist = 0;
	ap = 0;
	for (int i=0; i<size; i++)
		parse_attribute(&alist,argv[i]);
	alist = g_slist_sort(alist,str_cmp);
	int n = g_slist_length(alist);
	cout <<"n attributes = " << n << endl;
	char** encAtts = (char**)malloc((n+1)*sizeof(char*));
	int i=0;
	for (ap = alist; ap; ap = ap->next)
		encAtts[i++] = (char*)ap->data;
	encAtts[i] = 0;
	return encAtts;
}

int initOwner(char** attUniverse, int size, int b, unsigned char** pubBytes, int*
pubLen){

	beta = b;
	char** atts = parseAttributeUniverse(attUniverse,size);
	kpabe_setup(&pub, &msk, atts);
	toByteArray(kpabe_pub_serialize(pub), pubBytes, pubLen);
	
	//init the map
	initDiscreteLogMap(5000);

	//init the beta random values 
	maskArray = (element_t*) malloc(b*sizeof(element_t));
	for (int i=0; i<b; i++){
		element_init_Zr(maskArray[i],pub->p);
		element_random(maskArray[i]);
	}
	return 0;
}

int initDiscreteLogMap(int max){
	element_t exp;
	element_init_Zr(exp,pub->p);
	element_t val;
	element_init_GT(val,pub->p);
	for (int i=0; i<max; i++){
		element_set_si(exp,i);
		element_pow_zn(val,pub->gt,exp);
		element_wrap val1;
		element_init_GT(val1.element,pub->p);
		element_set(val1.element,val);
		mapElement[val1] = i; 	
		mapPlaintext[i] = val1;
	}
 	return 0;	
}

int toPlaintext_simple(element_t t){
	element_wrap et;
	element_init_GT(et.element,pub->p);

	element_set(et.element,t);

	unordered_map<element_wrap, int, elementHash, equalElement>::const_iterator it =
	mapElement.find(et);
		
	if (it==mapElement.end()){
		return -1;
	}
	return mapElement[et];

}

int toPlaintext_slidingWindow(element_t t, element_t mask){
	element_wrap et;
	element_init_GT(et.element,pub->p);

	element_t unmask, pt;
	element_init_GT(unmask, pub->p);
	element_init_GT(pt,pub->p);
	element_pow_zn(unmask,pub->gt,mask);
	element_div(pt,t,unmask);
	
	element_set(et.element,pt);

	unordered_map<element_wrap, int, elementHash, equalElement>::const_iterator it =
	mapElement.find(et);
		
	if (it==mapElement.end()){
		return -1;
	}
	return mapElement[et];

}

GByteArray* convertToByteArray(unsigned char* cipher, int len){
	GByteArray* buf = g_byte_array_new();
	g_byte_array_set_size(buf, len);
	for (int i=0; i<len; i++){
		buf->data[i] = cipher[i];
	}		
	return buf;
}

//encryption without masking.
int enc_simple(char** atts, int nAtts, int value, unsigned char** cipher, int* len){
	//parse the encryption attributes
	char** encAtts = parseEncryptionAttributes(atts,nAtts);


	//maping plaintext = gt^value
	element_wrap pt = mapPlaintext[value];;
	kpabe_cph_t *cph = kpabe_enc_raw(pub,pt.element,encAtts);

	GByteArray* buf = kpabe_cph_serialize(cph);

	*cipher = buf->data;
	*len = buf->len;
	return 0;
}

JNIEXPORT jbyteArray JNICALL Java_pCloud_crypto_kpabe_KPAbe_enc_1simple
  (JNIEnv *env, jobject obj, jobjectArray encAtts, jint val){
	int nAtts = env->GetArrayLength(encAtts);
	char** atts = (char**)malloc(nAtts*sizeof(char*));
	for (int i=0; i<nAtts; i++){
		atts[i] =
(char*)env->GetStringUTFChars((jstring)env->GetObjectArrayElement(encAtts, i), NULL);
	}	
	
	unsigned char* cipherBytes;
	int cipherLen;
	enc_simple(atts,nAtts, val, &cipherBytes, &cipherLen);
	jbyteArray ret = env->NewByteArray(cipherLen);
	env->SetByteArrayRegion(ret, 0, cipherLen, (jbyte*) cipherBytes);
	return ret;
}

//extractin Ep from kpabe_cph_t
JNIEXPORT jbyteArray JNICALL Java_pCloud_crypto_kpabe_KPAbe_extractOriginalCipher
  (JNIEnv *env, jobject obj, jbyteArray cipher){
	int cl = env->GetArrayLength(cipher);
	unsigned char *cipherBytes = (unsigned char*)malloc(cl*sizeof(char));
	env->GetByteArrayRegion(cipher,0,cl,(jbyte*)cipherBytes);
	GByteArray *buf = convertToByteArray(cipherBytes,cl);
	kpabe_cph_t *cph = kpabe_cph_unserialize(pub,buf,1);

	int n = element_length_in_bytes(cph->Ep);
	unsigned char* ep = (unsigned char*)malloc(n*sizeof(char));
	element_to_bytes(ep,cph->Ep);
	
	jbyteArray ret = env->NewByteArray(n);
	env->SetByteArrayRegion(ret,0,n,(jbyte*)ep);
	return ret;
}


int enc_slidingWindow(char** atts, int nAtts, int value, int index, unsigned char** cipher, int* len){
	//parse the encryption attributes
	char** encAtts = parseEncryptionAttributes(atts,nAtts);
	
	int factor = index/beta;
	
	element_t base2, exp, tmp, maskVal, ve, pt;
	element_init_Zr(base2, pub->p);
	element_init_Zr(exp, pub->p);
	element_set_si(base2, 2);
	element_set_si(exp, factor);
	
	element_init_Zr(tmp, pub->p);
	element_init_Zr(maskVal, pub->p);	
	
	element_pow_zn(tmp,base2,exp);
	element_mul(maskVal, tmp, maskArray[index%beta]);
	
	element_init_Zr(ve, pub->p);
	element_set_si(ve,value);
	
	element_add(tmp,ve,maskVal);	

	//here, tmp is the exponent value, pt is the plaintext mapped to GT
	element_init_GT(pt, pub->p);
	element_pow_zn(pt,pub->gt,tmp);

	kpabe_cph_t *cph = kpabe_enc_raw(pub,pt,encAtts);
	GByteArray* buf = kpabe_cph_serialize(cph);
	*cipher = buf->data;
	*len = buf->len;
	return 0;
}

JNIEXPORT jbyteArray JNICALL Java_pCloud_crypto_kpabe_KPAbe_enc_1slidingWindow
  (JNIEnv *env, jobject obj, jobjectArray encAtts, jint val, jint index){
	int nAtts = env->GetArrayLength(encAtts);
	char** atts = (char**)malloc(nAtts*sizeof(char*));
	for (int i=0; i<nAtts; i++){
		atts[i] =
(char*)env->GetStringUTFChars((jstring)env->GetObjectArrayElement(encAtts, i), NULL);
	}	
	
	unsigned char* cipherBytes;
	int cipherLen;
	enc_slidingWindow(atts,nAtts, val, index, &cipherBytes, &cipherLen);
	
	jbyteArray ret = env->NewByteArray(cipherLen);
	env->SetByteArrayRegion(ret, 0, cipherLen, (jbyte*) cipherBytes);
	return ret;
}

int generateTransformKey(char* policy, unsigned char** key, int policyId, int* len){
	cout <<" original policy = " << policy<<endl;
	char* parsedPolicy = parse_policy_lang(policy);
	element_t decK; 
	cout << "policy = " << parsedPolicy << endl;
	element_init_Zr(decK, pub->p);
	element_from_bytes(decK, userPolicyMap[policyId]);
	
	kpabe_prv_t* prvK = kpabe_keygen_proxy(pub,msk,parsedPolicy,decK);
	GByteArray* buf = kpabe_prv_serialize(prvK);
	*key = buf->data;
	*len = buf->len;
	return 0;
}
JNIEXPORT jbyteArray JNICALL Java_pCloud_crypto_kpabe_KPAbe_generateTransformKey
  (JNIEnv *env, jobject obj, jstring policy, jint id){
	char* p = (char*)env->GetStringUTFChars(policy, NULL);
	
	unsigned char* keyBytes;
	int keyLen;
	generateTransformKey(p, &keyBytes, id, &keyLen);
	jbyteArray ret = env->NewByteArray(keyLen);
	env->SetByteArrayRegion(ret,0,keyLen,(jbyte*)keyBytes);

	return ret;	
}

int generateUserDecryptionKey(unsigned char** key, int* len, int policyId){
	element_t decK;
	element_init_Zr(decK,pub->p);
	element_random(decK);
	*key = (unsigned char*)malloc(element_length_in_bytes(decK)*sizeof(char));		
	element_to_bytes(*key, decK);
	*len = element_length_in_bytes(decK);
	userPolicyMap[policyId] = *key;
	return 0;
}

JNIEXPORT jbyteArray JNICALL Java_pCloud_crypto_kpabe_KPAbe_generateUserDecryptionKey
  (JNIEnv *env, jobject obj, jint policyId){
	unsigned char *keyBytes;
	int keyLen;

	generateUserDecryptionKey(&keyBytes, &keyLen, policyId);
	jbyteArray ret = env->NewByteArray(keyLen);
	env->SetByteArrayRegion(ret,0,keyLen,(jbyte*)keyBytes);
	return ret;
}

int generateUserMaskValue(int alpha, unsigned char** mask, int* len){
	element_t base2, exp, tmp, sum, val, zero;
	element_init_Zr(base2, pub->p);
	element_init_Zr(exp, pub->p);
	element_init_Zr(tmp, pub->p);
	element_init_Zr(sum, pub->p);
	element_init_Zr(val, pub->p);
	element_init_Zr(zero, pub->p);
	
	element_set_si(base2, 2);
	element_set_si(zero,0);
	for (int i=0; i<beta; i++){
		element_set_si(exp, (alpha+i)/beta);
		element_pow_zn(tmp, base2, exp);
		element_mul(val, tmp, maskArray[(alpha+i)%beta]);
		element_add(sum, zero, val);
		element_set(zero,sum);
	}

	*len = element_length_in_bytes(sum);
	*mask = (unsigned char*)malloc((*len)*sizeof(char));
	element_to_bytes(*mask,sum);
}

JNIEXPORT jbyteArray JNICALL Java_pCloud_crypto_kpabe_KPAbe_generateUserMaskValue
  (JNIEnv *env, jobject obj, jint alpha){
	unsigned char* maskBytes;
	int maskLen;
	generateUserMaskValue(alpha, &maskBytes, &maskLen);
	jbyteArray ret = env->NewByteArray(maskLen);
	env->SetByteArrayRegion(ret,0,maskLen,(jbyte*)maskBytes);
	return ret;
}


JNIEXPORT jbyteArray JNICALL Java_pCloud_crypto_kpabe_KPAbe_initOwner
  (JNIEnv *env, jobject obj, jobjectArray atts, jint beta){
	int nAtts = env->GetArrayLength(atts);
	char** attsUniverse = (char**)malloc(nAtts*sizeof(char*));
	for (int i=0; i<nAtts; i++){
		attsUniverse[i] = (char*)env->GetStringUTFChars((jstring)env->GetObjectArrayElement(atts, i), NULL);
	}	
	
	unsigned char* pubBytes;
	int pubLen;
	initOwner(attsUniverse,nAtts, beta, &pubBytes, &pubLen);
	jbyteArray ret = env->NewByteArray(pubLen);
	env->SetByteArrayRegion(ret, 0, pubLen, (jbyte*) pubBytes);
	return ret;
}


